package com.wenyl.tree;

import lombok.Data;

@Data
public class RedBlackTree {
    private RBTNode root;
    public void insert(int value){
        // 新节点默认红色
        RBTNode newNode = new RBTNode(value);
        // 根节点为空，直接将新节点设置为根节点，颜色为黑色
        if(root == null){
            newNode.setColor(Color.BLACK);
            root = newNode;
            return;
        }

        RBTNode current = root;
        while(true){
            if(value < current.getValue()){
                if(current.getLeftChild() == null){
                    break;
                }
                current = current.getLeftChild();
            }else if(value > current.getValue()){
                if(current.getRightChild() == null){
                    break;
                }
                current = current.getRightChild();
            }else{
                System.out.println("节点已存在");
                return;
            }
        }
        if(newNode.getValue()< current.getValue()){
            current.setLeftChild(newNode);
        }
        if(newNode.getValue()>current.getValue()){
            current.setRightChild(newNode);
        }
        newNode.setFather(current);
        root = balanceTree(root,newNode);
    }

    private RBTNode balanceTree(RBTNode root,RBTNode current) {
        for(RBTNode father,grandFather,grandFatherLeftChild,grandFatherRightChild;;){
            father = current.getFather();
            if(father == null){
                current.setColor(Color.BLACK);
                return current;
            }
            // 如果父节点parent不是红色，或者父节点的父节点grandFather为空（即已经到达树的顶层），
            // 说明当前插入操作没有破坏红黑树的红黑性质（因为红节点不能连续出现两层，且根节点为黑色），
            // 此时无需进行进一步的平衡调整。
            grandFather = father.getFather();
            if(father.getColor() != Color.RED || grandFather == null){
                return root;
            }
            // 如果当前节点的父节点是爷爷节点的左孩子，则根据爷爷节点右孩子颜色判断
            grandFatherLeftChild = grandFather.getLeftChild();
            grandFatherRightChild = grandFather.getRightChild();
            if(father == grandFatherLeftChild){
                // 父亲是红色，叔叔节点是红色，则改变爷爷、父亲、叔叔节点的颜色，然后将current指向爷爷在判断是否平衡
                if(grandFatherRightChild != null && grandFatherRightChild.getColor() == Color.RED){
                    // 这里这么变色是因为叔叔是红色，那爷爷一定是黑色,所以爷爷变红色，两个孩子变黑色
                    grandFather.setColor(Color.RED);
                    grandFatherRightChild.setColor(Color.BLACK);
                    grandFatherLeftChild.setColor(Color.BLACK);
                    current = grandFather;
                }
                // 叔叔是黑色，说明不平衡情况不是上述的连续红色节点，需要进行旋转操作来调整树的结构
                else{
                    // 这里做了一个巧妙判断，先判断如果是LR型，则先对父节点进行左旋，变形为LL型号，在对爷爷节点右旋，这里变色是对爷爷和插入节点变色
                    // 在对爷爷进行右旋，但是入如果是LL型，则不会触发父节点的左旋，直接对爷爷进行右旋
                    // LL型是对爷爷和父亲节点变色，LR型是对当前节点和爷爷节点变色
                    if(current == father.getRightChild()){
                        current = father;
                        root = rotateLeft(root, current);
                        // 这里赋值后，插入节点变成了parent，然后更新爷爷节点，方便后面的变色和爷爷右旋操作
                        father = current.getFather();
                        grandFather = father == null ? null:father.getFather();
                    }
                    if(father != null){
                        father.setColor(Color.BLACK);
                        if(grandFather != null){
                            grandFather.setColor(Color.RED);
                            // 把爷爷右旋
                            root = rotateRight(root,grandFather);
                        }
                    }
                }
            }else{
                // 父亲是红色，叔叔节点是红色，则改变爷爷、父亲、叔叔节点的颜色，然后将current指向爷爷在判断是否平衡
                if(grandFatherLeftChild != null && grandFatherLeftChild.getColor() == Color.RED){
                    // 这里这么变色是因为叔叔是红色，那爷爷一定是黑色,所以爷爷变红色，两个孩子变黑色
                    grandFather.setColor(Color.RED);
                    grandFatherRightChild.setColor(Color.BLACK);
                    grandFatherLeftChild.setColor(Color.BLACK);
                    current = grandFather;
                }
                // 叔叔是黑色，说明不平衡情况不是上述的连续红色节点，需要进行旋转操作来调整树的结构
                else{
                    // 这里做了一个巧妙判断，先判断如果是RL型，则先对父节点进行右旋，变形为RR型，在对爷爷节点左旋，这里变色是对爷爷和插入节点变色
                    // 在对爷爷进行右旋，但是入如果是LL型，则不会触发父节点的左旋，直接对爷爷进行右旋
                    // LL型是对爷爷和父亲节点变色，LR型是对当前节点和爷爷节点变色
                    if(current == father.getLeftChild()){
                        current = father;
                        root = rotateRight(root, current);
                        // 这里赋值后，插入节点变成了parent，然后更新爷爷节点，方便后面的变色和爷爷右旋操作
                        father = current.getFather();
                        grandFather = father == null ? null:father.getFather();
                    }
                    if(father != null){
                        father.setColor(Color.BLACK);
                        if(grandFather != null){
                            grandFather.setColor(Color.RED);
                            // 把爷爷右旋
                            root = rotateLeft(root,grandFather);
                        }
                    }
                }
            }
        }
    }

    private RBTNode rotateRight(RBTNode root, RBTNode node) {
        RBTNode l,pp,lr;
        l = node.getLeftChild();
        if(node != null && l != null){
            lr = l.getRightChild();
            node.setLeftChild(lr);
            if(lr != null){
                lr.setFather(node);
            }
            pp = node.getFather();
            l.setFather(pp);
            if(pp == null){
                root = l;
                l.setColor(Color.BLACK);
            }else if(pp.getLeftChild() == node){
                pp.setLeftChild(l);
            }else{
                pp.setRightChild(l);
            }
            l.setRightChild(node);
            node.setFather(l);
        }
        return root;
    }

    // 把father节点左旋为右子节点的左孩子，右子节点的左孩子变为father的又孩子,然后把爷爷节点和father的右孩子链接
    private RBTNode rotateLeft(RBTNode root,RBTNode node){
        RBTNode r, pp, rl;
        r = node.getRightChild();
        if(node != null && r != null){
            rl = r.getLeftChild();
            node.setRightChild(rl);
            if(rl != null){
                rl.setFather(node);
            }
            pp = node.getFather();
            r.setFather(pp);
            if(pp == null){
                root = r;
                r.setColor(Color.BLACK);
            }else if(pp.getLeftChild() == node){
                pp.setLeftChild(r);
            }else{
                pp.setRightChild(r);
            }
            r.setLeftChild(node);
            node.setFather(r);
        }
        return root;
    }
}
enum Color{
    RED,
    BLACK;
}
@Data
class RBTNode{
    private int value;
    private RBTNode leftChild;
    private RBTNode rightChild;
    private RBTNode father;
    // 默认红色
    private Color color;
    // 默认为红色，保证根节点到叶子节点路径上的黑色节点数量相同
    public RBTNode(int value){
        this.value = value;
        this.color = Color.RED;
    }
    public void changeColor(){
        this.color = (this.color == Color.RED ? Color.BLACK : Color.RED);
    }
    @Override
    public String toString(){
        return "值="+value+" 颜色="+color+"\n";
    }
}       
